"""
根据指定输入的语言和内容做翻译,每次翻译一种言语。
1.建立大模型对象
2.创建消息提示词模板
3.创建基本的翻译链，
4.创建分支链执行的函数，不同语言分支链调用的函数
5.创建可运行分支链
6.执行分支链，输出结果
"""
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableBranch,RunnableLambda
from model_utils import getLLM

llm = getLLM()

#创建消息提示词模板
template = ChatPromptTemplate.from_messages([
    ("system","你是一个优秀的翻译官，将用户的信息翻译成{language}"),
    ("human","{content}")
])

parser = StrOutputParser()
#创建基本的翻译链
chain = template | llm | parser

def translate_fn(x):
    return RunnableLambda(lambda x : chain.invoke({"language":x["language"],"content":x["content"]}))

branch_chain = RunnableBranch(
    (lambda x : x["language"] == "英语",lambda x : translate_fn(x)),
    (lambda x : x["language"] == "法语",lambda x : translate_fn(x)),
    (lambda x : x["language"] == "德语",lambda x : translate_fn(x)),
    lambda x: f"{x['language']}不支持"
)

r = branch_chain.invoke({"language":"英语","content":"我爱你"})
print(r)